import torch
import torch.nn as nn

class ChannelAwareLoss(nn.Module):
    def __init__(self, snr_threshold=15):
        super().__init__()
        self.snr_threshold = snr_threshold
        self.ber_loss = nn.BCELoss()

    def forward(self, transmitted, received, H):
        """
        修正后的输入维度:
            transmitted: [B*K, Nc, 2Q] (实部虚部分开)
            received:    [B*K, Nc, 2Q]
            H:          [B, K, Nc, Nt] (复数形式，原始信道矩阵)
        """
        B_K, Nc, _ = transmitted.shape
        B = H.shape[0]
        K = B_K // B

        # 误码率计算 (保持原逻辑)
        tx_bits = (transmitted > 0).float()
        rx_bits = (received > 0).float()
        ber = self.ber_loss(rx_bits, tx_bits)

        # 信号功率约束 (维度调整)
        power = torch.mean(torch.abs(transmitted)**2, dim=(1,2))  # [B*K]

        # 信道功率计算 (关键修正)
        # 扩展H矩阵到每个用户
        H_expanded = H.squeeze()#.repeat(1, K, 1, 1, 1)  # [B, K, Nc, Nt]
        H_combined = H_expanded.reshape(B*K, Nc, -1)  # [B*K, Nc, Nt]

        # 计算等效SNR
        H_power = torch.mean(torch.abs(H_combined)**2, dim=(1,2))  # [B*K]
        snr_linear = H_power * power.view(B*K)  # 显式维度匹配
        snr_db = 10 * torch.log10(snr_linear + 1e-7)  # [B*K]

        # SNR自适应加权
        snr_weight = torch.sigmoid(snr_db - self.snr_threshold)
        power_penalty = torch.mean(power)

        return 0.7*ber + 0.3*power_penalty * torch.mean(snr_weight)